package cn.longwingstech.intelligence.agent.langchain4j;

import cn.longwingstech.intelligence.agent.monitor.AiModelMetricsCollector;
import dev.langchain4j.model.chat.listener.ChatModelErrorContext;
import dev.langchain4j.model.chat.listener.ChatModelListener;
import dev.langchain4j.model.chat.listener.ChatModelRequestContext;
import dev.langchain4j.model.chat.listener.ChatModelResponseContext;
import dev.langchain4j.model.output.TokenUsage;
import jakarta.annotation.Resource;
import org.springframework.stereotype.Component;

@Component
public class SuperAgentChatModelListener implements ChatModelListener {
    @Resource
    private AiModelMetricsCollector aiModelMetricsCollector;

    @Override
    public void onResponse(ChatModelResponseContext responseContext) {
        TokenUsage tokenUsage = responseContext.chatResponse().tokenUsage();
        Integer inputs = tokenUsage.inputTokenCount();
        Integer tokens = tokenUsage.outputTokenCount();
        Integer totals = tokenUsage.totalTokenCount();
        String modelName = responseContext.chatResponse().modelName();
        aiModelMetricsCollector.recordAllModelToken(modelName, "input", inputs);
        aiModelMetricsCollector.recordAllModelToken(modelName, "output", tokens);
        aiModelMetricsCollector.recordAllModelToken(modelName, "total", totals);
    }

}
